#include <jni.h>
#include <string>
#include <math.h>

#ifdef ANDROID
#include <jni.h>
#include <android/log.h>
#define LOGE(format, ...)  __android_log_print(ANDROID_LOG_ERROR, "xohnffmpeg", format, ##__VA_ARGS__)
#define LOGI(format, ...)  __android_log_print(ANDROID_LOG_INFO,  "xohnffmpeg", format, ##__VA_ARGS__)
#else
#define LOGE(format, ...)  printf("xohnffmpeg" format "\n", ##__VA_ARGS__)
#define LOGI(format, ...)  printf("xohnffmpeg " format "\n", ##__VA_ARGS__)
#endif


/**
 * Split Y, U, V planes in YUV420P file.
 * @param jstr_url Location of Input YUV file.
 * @param w    Width of Input YUV file.
 * @param h    Height of Input YUV file.
 * @param num  Number of frames to process.
 * 分离YUV420P像素数据中的Y、U、V分量
 * 如果视频帧的宽和高分别为w和h，那么一帧YUV420P像素数据一共占用w*h*3/2 Byte的数据。其中前w*h Byte存储Y，接着的w*h*1/4 Byte存储U，最后w*h*1/4 Byte存储V
 */
extern "C" JNIEXPORT jint JNICALL
Java_com_xohn_ffmpeg_AVUtils_yuv420Split(
        JNIEnv *env,jobject /* this */,
        jstring jstr_url, jint w, jint h, jint num) {

    char input_url[100] = {0};
    const char *str = env->GetStringUTFChars(jstr_url, NULL);
    sprintf(input_url,"%s",str);
    env->ReleaseStringUTFChars(jstr_url,str);

    FILE *fp = fopen(input_url,"rb+");
    if(fp == NULL) return -1;

    strcpy(input_url+strlen(input_url)-4,"_y.y");
    FILE *fp1 = fopen(input_url,"wb+"); //output_420_y.y
    strcpy(input_url+strlen(input_url)-4,"_u.y");
    FILE *fp2 = fopen(input_url,"wb+"); //output_420_u.y
    strcpy(input_url+strlen(input_url)-4,"_v.y");
    FILE *fp3 = fopen(input_url,"wb+"); //output_420_v.y

    unsigned char *pic = (unsigned char *)malloc(w*h*3/2);

    for(int i = 0;i < num;i++) {
        fread(pic, 1, w * h * 3 / 2, fp);
        //Y
        fwrite(pic, 1, w * h, fp1);
        //U
        fwrite(pic + w * h, 1, w * h / 4, fp2);
        //V
        fwrite(pic + w * h * 5 / 4, 1, w * h / 4, fp3);
    }

    free(pic);
    fclose(fp);
    fclose(fp1);
    fclose(fp2);
    fclose(fp3);

    return 0;
}


/**
 * Split Y, U, V planes in YUV444P file.
 * @param jstr_url Location of YUV file.
 * @param w    Width of Input YUV file.
 * @param h    Height of Input YUV file.
 * @param num  Number of frames to process.
 * 分离YUV444P像素数据中的Y、U、V分量
 * 如果视频帧的宽和高分别为w和h，那么一帧YUV444P像素数据一共占用w*h*3 Byte的数据。其中前w*h Byte存储Y，接着的w*h Byte存储U，最后w*h Byte存储V
 */
extern "C" JNIEXPORT jint JNICALL
Java_com_xohn_ffmpeg_AVUtils_yuv444Split(
        JNIEnv *env,jobject /* this */,
        jstring jstr_url, jint w, jint h, jint num) {
    char input_url[100] = {0};
    const char *str = env->GetStringUTFChars(jstr_url, NULL);
    sprintf(input_url,"%s",str);
    env->ReleaseStringUTFChars(jstr_url,str);

    FILE *fp = fopen(input_url,"rb+");
    if(fp == NULL) return -1;

    strcpy(input_url+strlen(input_url)-4,"_y.y");
    FILE *fp1 = fopen(input_url,"wb+");
    strcpy(input_url+strlen(input_url)-4,"_u.y");
    FILE *fp2 = fopen(input_url,"wb+");
    strcpy(input_url+strlen(input_url)-4,"_v.y");
    FILE *fp3 = fopen(input_url,"wb+");

    unsigned char *pic = (unsigned char *) malloc(w * h * 3);

    for(int i = 0;i < num;i++) {
        fread(pic, 1, w * h * 3, fp);
        //Y
        fwrite(pic, 1, w * h, fp1);
        //U
        fwrite(pic + w * h, 1, w * h, fp2);
        //V
        fwrite(pic + w * h * 2, 1, w * h, fp3);
    }

    free(pic);
    fclose(fp);
    fclose(fp1);
    fclose(fp2);
    fclose(fp3);

    return 0;
}


/**
 * Convert YUV420P file to gray picture
 * @param jstr_url Location of Input YUV file.
 * @param w       Width of Input YUV file.
 * @param h       Height of Input YUV file.
 * @param num     Number of frames to process.
 * 将YUV420P像素数据去掉颜色（变成灰度图）
 * 如果想把YUV格式像素数据变成灰度图像，只需要将U、V分量设置成128即可。
 * 这是因为U、V是图像中的经过偏置处理的色度分量。色度分量在偏置处理前的取值范围是-128至127，这时候的无色对应的是0值。经过偏置后色度分量取值变成了0至255，因而此时的无色对应的就是128了
 */
extern "C" JNIEXPORT jint JNICALL
Java_com_xohn_ffmpeg_AVUtils_yuv420Gray(
        JNIEnv *env,jobject /* this */,
        jstring jstr_url, jint w, jint h, jint num) {
    char input_url[100] = {0};
    const char *str = env->GetStringUTFChars(jstr_url, NULL);
    sprintf(input_url,"%s",str);
    env->ReleaseStringUTFChars(jstr_url,str);

    FILE *fp = fopen(input_url,"rb+");
    if(fp == NULL) return -1;

    strcpy(input_url+strlen(input_url)-4,"_gray.yuv");
    FILE *fp1 = fopen(input_url,"wb+");

    unsigned char *pic=(unsigned char *)malloc(w*h*3/2);

    for(int i = 0;i < num;i++) {
        fread(pic, 1, w * h * 3 / 2, fp);
        //Gray
        memset(pic + w * h, 128, w * h / 2);
        fwrite(pic, 1, w * h * 3 / 2, fp1);
    }

    free(pic);
    fclose(fp);
    fclose(fp1);
    return 0;
}


/**
 * Halve Y value of YUV420P file
 * @param jstr_url Location of Input YUV file.
 * @param w       Width of Input YUV file.
 * @param h       Height of Input YUV file.
 * @param num     Number of frames to process.
 * 将YUV420P像素数据的亮度减半
 * 如果打算将图像的亮度减半，只要将图像的每个像素的Y值取出来分别进行除以2就可以了
 */
extern "C" JNIEXPORT jint JNICALL
Java_com_xohn_ffmpeg_AVUtils_yuv420Halfy(
        JNIEnv *env,jobject /* this */,
        jstring jstr_url, jint w, jint h, jint num) {
    char input_url[100] = {0};
    const char *str = env->GetStringUTFChars(jstr_url, NULL);
    sprintf(input_url,"%s",str);
    env->ReleaseStringUTFChars(jstr_url,str);

    FILE *fp = fopen(input_url,"rb+");
    if(fp == NULL) return -1;

    strcpy(input_url+strlen(input_url)-4,"_halfy.yuv");
    FILE *fp1 = fopen(input_url,"wb+");

    unsigned char *pic=(unsigned char *)malloc(w*h*3/2);

    for(int i = 0;i < num;i++) {
        fread(pic, 1, w * h * 3 / 2, fp);
        //Half
        for (int j = 0; j < w * h; j++) {
            unsigned char temp = pic[j] / 2;
            pic[j] = temp;
        }
        fwrite(pic, 1, w * h * 3 / 2, fp1);
    }

    free(pic);
    fclose(fp);
    fclose(fp1);

    return 0;
}


/**
 * Add border for YUV420P file
 * @param jstr_url Location of Input YUV file.
 * @param w       Width of Input YUV file.
 * @param h       Height of Input YUV file.
 * @param border  Width of Border.
 * @param num     Number of frames to process.
 * 将YUV420P像素数据的周围加上边框
 * 图像的边框的宽度为border，本程序将距离图像边缘border范围内的像素的亮度分量Y的取值设置成了亮度最大值255
 */
extern "C" JNIEXPORT jint JNICALL
Java_com_xohn_ffmpeg_AVUtils_yuv420Border(
        JNIEnv *env,jobject /* this */,
        jstring jstr_url, jint w, jint h, jint border, jint num) {
    char input_url[100] = {0};
    const char *str = env->GetStringUTFChars(jstr_url, NULL);
    sprintf(input_url,"%s",str);
    env->ReleaseStringUTFChars(jstr_url,str);

    FILE *fp = fopen(input_url,"rb+");
    if(fp == NULL) return -1;

    strcpy(input_url+strlen(input_url)-4,"_border.yuv");
    FILE *fp1 = fopen(input_url,"wb+");

    unsigned char *pic=(unsigned char *)malloc(w*h*3/2);

    for(int i = 0;i < num;i++){
        fread(pic,1,w*h*3/2,fp);
        //Y
        for(int j = 0;j < h;j++){
            for(int k = 0;k < w;k++){
                if(k < border || k >= (w-border) || j < border || j >= (h-border)){
                    pic[j*w+k]=255;
                    //pic[j*w+k]=0;
                }
            }
        }
        fwrite(pic,1,w*h*3/2,fp1);
    }

    free(pic);
    fclose(fp);
    fclose(fp1);

    return 0;
}


/**
 * Generate YUV420P gray scale bar.
 * @param width    Width of Output YUV file.
 * @param height   Height of Output YUV file.
 * @param ymin     Max value of Y
 * @param ymax     Min value of Y
 * @param barnum   Number of bars
 * @param jstr_url_out Location of Output YUV file.
 * 生成YUV420P格式的灰阶测试图
 * 通过灰阶测试图的亮度最小值ymin，亮度最大值ymax，灰阶数量barnum确定每一个灰度条中像素的亮度分量Y的取值。
 * 还要根据图像的宽度width和图像的高度height以及灰阶数量barnum确定每一个灰度条的宽度
 */
extern "C" JNIEXPORT jint JNICALL
Java_com_xohn_ffmpeg_AVUtils_yuv420GrayBar(
        JNIEnv *env,jobject /* this */,
        jstring jstr_url_out, jint width, jint height, jint ymin,jint ymax,jint barnum) {
    int barwidth;
    float lum_inc;
    unsigned char lum_temp;
    int uv_width,uv_height;
    FILE *fp=NULL;
    unsigned char *data_y = NULL;
    unsigned char *data_u = NULL;
    unsigned char *data_v = NULL;
    int t = 0,i = 0,j = 0;

    barwidth = width/barnum;
    lum_inc = ((float)(ymax-ymin))/((float)(barnum-1));
    uv_width = width/2;
    uv_height = height/2;

    data_y = (unsigned char *)malloc(width*height);
    data_u = (unsigned char *)malloc(uv_width*uv_height);
    data_v = (unsigned char *)malloc(uv_width*uv_height);

    char input_url[100] = {0};
    const char *str = env->GetStringUTFChars(jstr_url_out, NULL);
    sprintf(input_url,"%s",str);
    env->ReleaseStringUTFChars(jstr_url_out,str);

    if((fp=fopen(input_url,"wb+")) == NULL){
        LOGE("Error: Cannot create file!");
        return -1;
    }

    //Output Info
    LOGI("Y, U, V value from picture's left to right:\n");
    for(t = 0;t < barnum;t++){
        lum_temp = ymin+(char)(t*lum_inc);
        LOGI("%3d, 128, 128\n",lum_temp);
    }
    //Gen Data
    for(j = 0;j < height;j++){
        for(i = 0;i < width;i++){
            t = i/barwidth;
            lum_temp = ymin+(char)(t*lum_inc);
            data_y[j*width+i] = lum_temp;
        }
    }
    for(j = 0;j < uv_height;j++){
        for(i = 0;i < uv_width;i++){
            data_u[j*uv_width+i] = 128;
        }
    }
    for(j = 0;j < uv_height;j++){
        for(i = 0;i < uv_width;i++){
            data_v[j*uv_width+i] = 128;
        }
    }
    fwrite(data_y,width*height,1,fp);
    fwrite(data_u,uv_width*uv_height,1,fp);
    fwrite(data_v,uv_width*uv_height,1,fp);
    fclose(fp);
    free(data_y);
    free(data_u);
    free(data_v);
    return 0;
}


/**
 * Calculate PSNR between 2 YUV420P file
 * @param url1     Location of first Input YUV file.
 * @param url2     Location of another Input YUV file.
 * @param w        Width of Input YUV file.
 * @param h        Height of Input YUV file.
 * @param num      Number of frames to process.
 * 计算两个YUV420P像素数据的PSNR
 * PSNR是最基本的视频质量评价方法
 * 程序计算后得到的PSNR取值为26.693。PSNR取值通常情况下都在20-50的范围内，取值越高，代表两张图像越接近，反映出受损图像质量越好
 */
extern "C" JNIEXPORT jint JNICALL
Java_com_xohn_ffmpeg_AVUtils_yuv420Psnr(
        JNIEnv *env,jobject /* this */,
        jstring jstr_url1,jstring jstr_url2, jint w, jint h,jint num) {
    const char *str1 = env->GetStringUTFChars(jstr_url1, NULL);
    FILE *fp1 = fopen(str1,"rb+");
    env->ReleaseStringUTFChars(jstr_url1,str1);
    if(fp1 == NULL)
        return  -1;
    const char *str2 = env->GetStringUTFChars(jstr_url2, NULL);
    FILE *fp2 = fopen(str2,"rb+");
    env->ReleaseStringUTFChars(jstr_url2,str2);
    if(fp2 == NULL)
        return  -1;

    unsigned char *pic1 = (unsigned char *)malloc(w*h);
    unsigned char *pic2 = (unsigned char *)malloc(w*h);

    for(int i = 0;i < num;i++){
        fread(pic1,1,w*h,fp1);
        fread(pic2,1,w*h,fp2);

        double mse_sum = 0,mse = 0,psnr = 0;
        for(int j = 0;j < w*h;j++){
            mse_sum += pow((double)(pic1[j]-pic2[j]),2);
        }
        mse = mse_sum/(w*h);
        psnr = 10*log10(255.0*255.0/mse);
        LOGI("%5.3f\n",psnr);

        fseek(fp1,w*h/2,SEEK_CUR);
        fseek(fp2,w*h/2,SEEK_CUR);

    }

    free(pic1);
    free(pic2);
    fclose(fp1);
    fclose(fp2);
    return 0;
}





extern "C" JNIEXPORT jstring JNICALL
Java_com_xohn_ffmpeg_AVUtils_stringFromJNI(
        JNIEnv *env,jobject /* this */) {
    std::string hello = "Hello FFmpeg from C++";
    return env->NewStringUTF(hello.c_str());
}
